package vip.meeet;

import org.apache.http.client.config.RequestConfig;
import org.apache.http.client.methods.CloseableHttpResponse;
import org.apache.http.client.methods.HttpPost;
import org.apache.http.entity.StringEntity;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.impl.client.HttpClients;
import org.apache.http.util.EntityUtils;

import java.io.IOException;

/**
 * elasticsearch bulk 请求
 * https://www.elastic.co/guide/en/elasticsearch/reference/7.4/docs-bulk.html
 */
public class BulkBuilder {

    private static final String INDEX_TEMPLATE_WITH_ID = "{ \"%s\" : { \"_index\" : \"%s\", \"_id\" : \"%s\" } }";

    private static final String INDEX_TEMPLATE = "{ \"%s\" : { \"_index\" : \"%s\"} }";

    private static final String DOC_TEMPLATE = "{ \"doc\" : %s }";

    private static final String INDEX = "index";
    private static final String CREATE = "create";
    private static final String DELETE = "delete";
    private static final String UPDATE = "update";

    /**
     * 换行、允许是\n或者\r\n
     */
    private static final String NEWLINE = "\n";

    /**
     * bulk请求的content-type
     */
    private static final String APPLICATION_X_NDJSON = "application/x-ndjson";

    private static final String UTF_8 = "UTF-8";

    String BULK_URL = "http://%s/_bulk";


    private StringBuffer body;

    private BulkBuilder(){
        body = new StringBuffer();
    }

    public static BulkBuilder getBuilder(){
        return new BulkBuilder();
    }

    public BulkBuilder index(String index, String id, String docJson){
        String head = String.format(INDEX_TEMPLATE_WITH_ID, INDEX, index, id);
        body.append(head).append(NEWLINE);
        body.append(docJson).append(NEWLINE);
        return this;
    }

    public BulkBuilder index(String index, String docJson){
        String head = String.format(INDEX_TEMPLATE, INDEX, index);
        body.append(head).append(NEWLINE);
        body.append(docJson).append(NEWLINE);
        return this;
    }

    public BulkBuilder update(String index, String id, String docJson){
        String head = String.format(INDEX_TEMPLATE_WITH_ID, UPDATE, index, id);
        body.append(head).append(NEWLINE);
        String updateBody = String.format(DOC_TEMPLATE, docJson);
        body.append(updateBody).append(NEWLINE);
        return this;
    }

    public BulkBuilder create(String index, String id, String docJson){
        String head = String.format(INDEX_TEMPLATE_WITH_ID, CREATE, index, id);
        body.append(head).append(NEWLINE);
        body.append(docJson).append(NEWLINE);
        return this;
    }

    public BulkBuilder create(String index, String docJson){
        String head = String.format(INDEX_TEMPLATE, CREATE, index);
        body.append(head).append(NEWLINE);
        body.append(docJson).append(NEWLINE);
        return this;
    }

    public BulkBuilder delete(String index, String id){
        String head = String.format(INDEX_TEMPLATE_WITH_ID, DELETE, index, id);
        body.append(head).append(NEWLINE);
        return this;
    }

    public String build(){
        return body.toString();
    }

    public String postBulk(String hostPort) throws IOException {
        String url = String.format(BULK_URL, hostPort);
        if(body.length() > 15){
            return bulkRequest(url,body.toString());
        }else {
            return "";
        }
    }

    private static String bulkRequest(String url,String body) throws IOException {
        CloseableHttpClient httpClient = HttpClients.createDefault();//使用http连接池遇到 no shard found问题
        HttpPost method = new HttpPost(url);
        RequestConfig.Builder custom = RequestConfig.custom();
        custom.setSocketTimeout(1000 * 5 * 60);//根据数据量调整，如果read socketTimeout
        method.setConfig(custom.build());
        method.setHeader("Charset",UTF_8);//根据实际编码调整
        method.setHeader("Content-Type", APPLICATION_X_NDJSON);
        StringEntity entity = new StringEntity(body, UTF_8);
        method.setEntity(entity);
        CloseableHttpResponse response = httpClient.execute(method);
        String rs = EntityUtils.toString(response.getEntity());
        httpClient.close();
        return rs;
    }
}
